package com.winner.data.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.stereotype.Component;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import javax.annotation.Resource;
import java.util.List;
import java.util.function.BiFunction;

/**
 * Mybatis批处理 数据库连接信息拼接 &rewriteBatchedStatements=true 开启批量
 * 数据库主键自增也是支持批量插入后返回主键的, 只需要在xml里使用
 * <insert id="insertOne" useGeneratedKeys="true" keyProperty="id">
 * </insert>
 *
 * MyBatis-Plus saveBatch保存太慢解决办法
 * https://blog.csdn.net/qq_27628011/article/details/127527926
 * 使用mybatisplus的saveBatch(list, 3000默认不写是1000)也是可以批量保存的, 前提是数据库url &rewriteBatchedStatements=true 开启批量, 因为和该工具类的原理是一样的
 *
 * 批量删除 根据索引字段删除 用 行行比较删除 和 批处理删除都行, 二者效率差不多; 不是根据索引字段删除,一定要用in, 200w数据删除200条用时7秒.  如果用批处理删除, 需要半小时
 *
 * @author wangxinle5
 * @since 2022-09-26
 */
@Slf4j
@Component
public class MybatisBatchUtil {


//    设置适当的批处理大小：批处理大小指在一次插入操作中插入多少行数据。
//    如果批处理值太小，插入操作的频率将很高，而如果批处理值太大，可能会导致内存占用过高。
//    通常，建议将批处理大小设置为1000-5000行，这将减少插入操作的频率并降低内存占用。
    /**
     * 每次批处理的条数
     */
    public static final int BATCH_SIZE = 3000;

    @Resource
    private SqlSessionFactory sqlSessionFactory;

    // 传入的list可以不用是数据库的 List<User>, 只要有insertOne(实体类映射user表的字段即可)
    // 需要自己写一个insertOne updateOne的mapper对应的sql, 注意resultType不写,但是mapper接口里必须返回值是int
    // int insertAffectRow = mybatisBatchUtil.batchUpdateOrInsert(list, UserMapper.class, (item, mapper) ->
    // mapper.insertOne(item));
    // int updateAffectRow = mybatisBatchUtil.batchUpdateOrInsert(list, UserMapper.class, (item, mapper) ->
    // mapper.updateOne(item));
    // int deleteAffectRow = mybatisBatchUtil.batchUpdateOrInsert(list, UserMapper.class, (item, mapper) ->
    // mapper.deleteOne(item));

    /**
     * 批量处理(新增/更新/删除 都可以用) 主键冲突会直接回滚   对于多线程异步保存更新数据, 可以将@Async注解加到业务实现类上
     *
     * <insert id="insertOne" useGeneratedKeys="true" keyProperty="userId"> 就支持批量保存返回自增的主键, 返回的主键就在入参dbList里 mybatisBatchUtil.batchUpdateOrInsert(dbList, UserMapper.class,
     *
     * @param data        需要被处理的数据
     * @param mapperClass Mybatis的Mapper类
     * @param function    自定义处理逻辑
     * @return int 影响的总行数
     */
    public <T, U, R> int batchUpdateOrInsert(List<T> data, Class<U> mapperClass, BiFunction<T, U, R> function) {
        int i = 1;
        SqlSession batchSqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
        try {
            U mapper = batchSqlSession.getMapper(mapperClass);
            int size = data.size();
            for (T element : data) {
                function.apply(element, mapper);
                // 每当 i 的值达到 BATCH_SIZE 的倍数，或者当 i 达到 data 列表的大小时会调用 batchSqlSession.flushStatements() 来执行批处理保存
                if ((i % 1000 == 0) || i == size) {
                    batchSqlSession.flushStatements();
                }
                i++;
            }
            // 非事务环境下强制commit，事务情况下该commit相当于无效
            batchSqlSession.commit(!TransactionSynchronizationManager.isSynchronizationActive());
        } catch (Exception e) {
            // 如果一批保存的数据量有主键冲突异常, 那么这批BATCH_SIZE大小的数据都不会入库.
            batchSqlSession.rollback();
            log.error("mybatis批处理报错, error is {}", e.getMessage());
            // throw new CustomException(e);
        } finally {
            batchSqlSession.close();
        }
        return i - 1;
    }

}
